from urllib.request import urlopen
import requests

import tidal.artist
import tidal.album
from generic.track import Track, guess_file_name
from generic.utils import print_progress_callback

class TidalTrack(Track):
  def set_raw_data(self, song_data):
    self.set_values(track_id = song_data["id"],
                    title = song_data["title"],
                    artists = list(map(
                      lambda a: tidal.artist.TidalArtist.grab_or_create(
                        a["id"], raw_data = a
                      ), song_data["artists"]
                    )),
                    album = tidal.album.TidalAlbum.grab_or_create(
                      song_data["album"]["id"],
                      album_id = song_data["album"]["id"],
                      title = song_data["album"]["title"],
                      is_full = False
                    ),
                    qualities = [song_data["audioQuality"]],
                    is_full = True)

  def set_values(self,
                 track_id = None,
                 title = None,
                 artists = None,
                 album = None,
                 qualities = None,
                 is_full = False, # if track is fully loaded
                 raw_data = None):
    if raw_data != None:
      self.set_raw_data(raw_data)
      return
    self.track_id = track_id
    self.title = title
    self.artists = artists
    self.album = album
    self.qualities = qualities
    self.is_full = is_full

  def get_full_data(self, tidal_instance):
    if not self.is_full:
      self.set_raw_data(tidal_instance.get_track_data(self.track_id))

  def get_stream(self, tidal_instance, quality):
    return tidal_instance.get_stream(self.track_id, quality)

  def stream(self, tidal_instance, output_stream, progress_callback = None):
    track_stream = self.get_stream(tidal_instance, quality)
    self.stream(track_stream,
                output_stream,
                progress_callback = progress_callback)
  def _stream(self, track_stream, output_stream, progress_callback = None):
    stream = urlopen(track_stream.url)

    length = int(stream.getheader('Content-Length'))
    chunk_size = 8192
    i = 0
    while True:
      chunk = stream.read(chunk_size)
      if progress_callback != None:
        if (i * chunk_size) > length:
          progress_callback(1)
        else:
          progress_callback((i * chunk_size) / length)
      if not chunk:
        break
      output_stream.write(chunk)
      i += 1

  def download(self,
               tidal_instance,
               to_file = None,
               quality = None,
               print_progress = False):
    if quality == None:
      quality = self.guess_quality()
    if to_file == None:
      to_file = guess_file_name(self)
    track_stream = self.get_stream(tidal_instance, quality)
    file_format = ".flac" if ".m4a" not in track_stream.url else ".m4a"
    file_name = to_file + file_format

    with open(file_name, 'wb') as stream:
      self._stream(track_stream,
                   stream, # output stream
                   progress_callback = print_progress_callback \
                   if print_progress else None)
